from typing import List

from pandas import Series
from sklearn.linear_model import LinearRegression
import pandas as pd
import re
import openai

GET_TRAIN_DATA_TEMPLATE = ""
GET_PREDICT_FEATURES_TEMPLATE = """用户发言被三个反引号括住。将用户发言中待预测的特征值格式化为Python的dict并输出，dict的value应当是浮点数，输出时反引号不要带编程语言名称。```{query}```"""

openai.api_base = "http://10.212.253.250:8082/v1"
openai.api_key = "none"


def get_predict_feature(query: str) -> Series | None:
    openai.api_base = "http://10.212.253.250:8082/v1"
    openai.api_key = "none"
    cp = openai.ChatCompletion.create(
        model="chatglm2-6b",
        messages=[
            {"role": "user",
             "content": GET_PREDICT_FEATURES_TEMPLATE.replace("query", query)}
        ]
    )
    res = cp.choices[0].message.content
    z = re.findall("```[^`]*```", res, re.DOTALL)
    print("z: ", z)
    for str_ in z[::-1]:
        d = str_[3:-3].strip()
        try:
            c = pd.Series(eval(d))
        except SyntaxError:
            continue
        else:
            print(c)
            return [c]
    else:
        return


def get_train_data(query: str) -> (List[List], List):
    return "", ""


def get_data(query: str, mode="t") -> (List[List] | None, List | None):
    features, labels = None, None
    if mode == 't':
        features, labels = get_train_data(query)
    elif mode == 'p':
        features = get_predict_feature(query)
        labels = None
    return features, labels


class LR:

    def __init__(self):
        self.active = False
        self.model = LinearRegression()
        self.model.fit(
            [[0, 0, 1], [1, 2, 2], [2, 3, 3]],
            [0, 1, 2]
        )
        self.features_ = None
        self.labels_ = None

    def reset(self):
        self.__init__()

    def append(self, x, y):
        self.features_ += x
        self.labels_ += y

    def deal(self, query: str, history: list):
        ans = ""
        if "线性回归" in query:
            self.active = True
            ans += "开始线性回归，请输入您的训练集~"
        if self.active:
            if "训练" in query:
                self.model.fit(*get_data(query, "t"))
                ans += f"训练完成~\n\t线性系数：{self.model.coef_}\n\t截距：{self.model.intercept_}\n可以进行预测啦~"
            if "预测" in query:
                features = get_data(query, "p")[0]
                if features is None:
                    ans += "输入不明确，请重新输入您的预测条件。"
                else:
                    res = self.model.predict(features)
                    ans += f"预测结果为：{res}\n"
            if "结束" in query:
                self.reset()
                ans += f"回归结束"
        history.append([query, ans])
        return history, ''
